import gradio as gr
import os
import argparse

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI

import category


def initialize_sales_bot(cate: category.Category, vector_store_dir: str = "sales_chatbot/vector_store/"):
    db = FAISS.load_local(vector_store_dir+cate.key, OpenAIEmbeddings())
    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

    global SALES_BOT
    SALES_BOT = RetrievalQA.from_chain_type(llm,
                                            retriever=db.as_retriever(search_type="similarity_score_threshold",
                                                                      search_kwargs={"score_threshold": 0.8}))
    # 返回向量数据库的检索结果
    SALES_BOT.return_source_documents = True

    return SALES_BOT


def sales_chat(message, history):
    print(f"[message]{message}")
    print(f"[history]{history}")

    ans = SALES_BOT({"query": message})
    # 如果检索出结果，或者开了大模型聊天模式
    # 返回 RetrievalQA combine_documents_chain 整合的结果
    if ans["source_documents"] or enable_chat:
        print(f"[result]{ans['result']}")
        print(f"[source_documents]{ans['source_documents']}")
        return ans["result"]
    # 否则输出套路话术
    else:
        return "这个问题我要问问领导"


def launch_gradio(cate: category.Category):
    demo = gr.ChatInterface(
        fn=sales_chat,
        title=cate.name+"销售",
        theme="soft",
        retry_btn=None,
        undo_btn=None,
        chatbot=gr.Chatbot(height=600),
    )

    demo.launch(share=True, server_name="0.0.0.0")


global enable_chat
enable_chat = False

if __name__ == "__main__":
    os.environ.setdefault(
        'OPENAI_API_KEY', 'XXX')

    parser = argparse.ArgumentParser()
    parser.add_argument("--name", default='房产', help='销售分类:' +
                        str([cate.name for cate in category.category]))
    parser.add_argument("--enable_chat", default=False, help='是否启用LLM:默认False')

    args = parser.parse_args()
    print('启动销售类型: ', args.name)
    enable_chat = args.enable_chat
    if enable_chat:
        print('启动LLM聊天')

    cate = category.default_category
    is_find = False
    for c in category.category:
        if c.name == args.name:
            cate = c
            is_find = True
            break

    if not is_find:
        print('未找到指定销售类型, 使用默认类型:', cate.name)

    # 初始化销售机器人
    initialize_sales_bot(cate)
    # 启动 Gradio 服务
    launch_gradio(cate)
